#include "kernel_operator.h"
// #include <limits>
using namespace AscendC;
constexpr int32_t BUFFER_NUM = 2;
// float positiveInfinity = std::numeric_limits<float>::infinity();
// float negativeInfinity = -std::numeric_limits<float>::infinity();



template<typename T>
class KernelLpnormV2{
public:
    __aicore__ inline KernelLpnormV2() {}
    __aicore__ inline void Init(GM_ADDR x, GM_ADDR y,
                            uint32_t blockLength,
                            uint32_t tileNum, uint32_t tileLength,
                            uint32_t lasttileLength, uint32_t formerNum,
                            uint32_t formerLength, uint32_t formertileNum,
                            uint32_t formertileLength,
                            uint32_t formerlasttileLength, uint32_t tailNum,
                            uint32_t tailLength, uint32_t tailtileNum,
                            uint32_t tailtileLength,
                            uint32_t taillasttileLength, uint32_t tilingKey, int size, float p) {
    this->length = size;
    this->p = p;
    ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");
    if (tilingKey == 1) {
        this->blockLength = blockLength;
        this->tileNum =
            tileNum ASSERT(tileNum != 0 && "tile num can not be zero!");
        this->tileLength = tileLength / BUFFER_NUM;
        this->lasttileLength = lasttileLength;

        xGm.SetGlobalBuffer((__gm__ T*)x + this->blockLength * GetBlockIdx(),
                            this->blockLength);               
        yGm.SetGlobalBuffer((__gm__ T*)y + this->blockLength * GetBlockIdx(),
            this->blockLength);
    }

    if (tilingKey == 2) {
        this->formerNum = formerNum;
        this->formerLength = formerLength;
        this->formertileNum = formertileNum;
        this->formertileLength = formertileLength;
        this->formerlasttileLength = formerlasttileLength;

        this->tailNum = tailNum;
        this->tailLength = tailLength;
        this->tailtileNum = tailtileNum;
        this->tailtileLength = tailtileLength;
        this->taillasttileLength = taillasttileLength;

        if (GetBlockIdx() < this->formerNum) {  //分到大块核的处理
            this->tileLength = this->formertileLength / BUFFER_NUM;
            this->lasttileLength = this->formerlasttileLength;
            this->tileNum = this->formertileNum * BUFFER_NUM;
            xGm.SetGlobalBuffer(
                (__gm__ T*)x + this->formerLength * GetBlockIdx(),
                this->formerLength);              
            yGm.SetGlobalBuffer(
                (__gm__ T*)y + this->formerLength * GetBlockIdx(),
                this->formerLength);
        } else {  //分到小块核的处理，需要处理的数据量比大核少alignNum个
            this->tileLength = this->tailtileLength / BUFFER_NUM;
            this->lasttileLength = this->taillasttileLength;
            this->tileNum = this->tailtileNum * BUFFER_NUM;
            xGm.SetGlobalBuffer(
                (__gm__ T*)x + this->formerLength * this->formerNum +
                    this->tailLength * (GetBlockIdx() - this->formerNum),
                this->tailLength);
            yGm.SetGlobalBuffer(
                (__gm__ T*)y + this->formerLength * this->formerNum +
                    this->tailLength * (GetBlockIdx() - this->formerNum),
                this->tailLength);
        }
    }

    pipe.InitBuffer(inQueueIN, BUFFER_NUM, this->tileLength * 1 * sizeof(T));
    pipe.InitBuffer(outQueueOUT, BUFFER_NUM, this->tileLength * sizeof(T));

    pipe.InitBuffer(inQueueC, this->tileLength * sizeof(T));
    pipe.InitBuffer(tmpQ, this->tileLength * sizeof(int32_t));
    }

    __aicore__ inline void Process(float p) {
        int32_t loopCount = this->tileNum * BUFFER_NUM;
        for (int32_t i = 0; i < loopCount; i++) {
            CopyIn(i);
            if( p == (float)2.0 )
            {
                Compute_defaut(i, p);
            }
            // else if(p == positiveInfinity)
            // {
            //     Compute_inf(i, p);
            // }else if(p == negativeInfinity)
            // {
            //     Compute_inf(i, p);
            // }
            else if(p == (float)0.0)
            {
                Compute_0(i, p);
            }else if(p == (float)1.0)
            {
                Compute_inf(i, p);
            }
            // else 
            // {
            //     Compute_else(i, p);
            // }
            // Compute(i, p);
            CopyOut(i);
        }
        pipe_barrier(PIPE_ALL);
        if( p == (float)2.0 ) 
        { 
            float value = 0;
            for(int i = 0; i < this->length; i++)
            {
                value += (float)yGm.GetValue(i);
            }
                        
            for (int32_t i = 0; i < loopCount; i++) 
            {
                FinalComput(i, p, value);
            }        
        }
        // else if(p == positiveInfinity) 
        // {
        //     float maxValue = -10000.0;
        //     // 求max
        //     for(int i = 0; i < this->length; i++)
        //     {
        //         if(yGm.GetValue(i) > maxValue)
        //         {
        //             maxValue = yGm.GetValue(i);
        //         }
        //     }
            
        //     for (int32_t i = 0; i < loopCount; i++) 
        //     {
        //         FinalComput(i, p, maxValue);
        //     }
        // }
        // else if(p == negativeInfinity) 
        // {
        //     float minValue = 100000.0;
        //     // 求max
        //     for(int i = 0; i < this->length; i++)
        //     {
        //         if(yGm.GetValue(i) < minValue)
        //         {
        //             minValue = yGm.GetValue(i);
        //         }
        //     }
            
        //     for (int32_t i = 0; i < loopCount; i++) 
        //     {
        //         FinalComput(i, p, minValue);
        //     }
        // }
        else if(p == (float)0.0)
        {
            float value = 0;
            for(int i = 0; i < this->length; i++)
            {
                value += (float)yGm.GetValue(i);
            }                     
            for (int32_t i = 0; i < loopCount; i++) 
            {
                FinalComput(i, p, value);
            }
        }
        else if(p == (float)1.0)
        {   
            float value = 0;
            for(int i = 0; i < this->length; i++)
            {
                value += (float)yGm.GetValue(i);
            }                     
            for (int32_t i = 0; i < loopCount; i++) 
            {
                FinalComput(i, p, value);
            }
        }
        // else
        // {
        //     float value = 0;
        //     for(int i = 0; i < this->length; i++)
        //     {
        //         value += (float)yGm.GetValue(i);
        //     }
        //     for(int i = 0)   
        // }

    }

private:
    __aicore__ inline void CopyIn(int32_t progress) {
        LocalTensor<T> inLocal = inQueueIN.AllocTensor<T>();
        if (BUFFER_NUM == 1) {
            if (progress == this->tileNum - 1) {
                if (progress == 0) {
                //如果只有一包，则搬运的起始地址为0，tileLength为实际分块的数据量
                DataCopy(inLocal[0], xGm[0], this->tileLength);
                DataCopy(inLocal[this->tileLength], labelGm[0], this->tileLength);                
                } else {
                //将最后一个分块的起始地址向前移动tileLength-lasttileLength
                DataCopy(
                    inLocal[0],
                    xGm[(progress - 1) * this->tileLength + this->lasttileLength],
                    this->tileLength);
                DataCopy(
                    inLocal[this->tileLength],
                    labelGm[(progress - 1) * this->tileLength + this->lasttileLength],
                    this->tileLength);                    
                }
            } else {
                DataCopy(inLocal[0], xGm[progress * this->tileLength],
                        this->tileLength);
                DataCopy(inLocal[this->tileLength], labelGm[progress * this->tileLength],
                        this->tileLength);                        
            }
        }
        if (BUFFER_NUM == 2) {
            //开启double
            //buffer时，由于将输入数据分成了相等的2部分，分块大小为不开启double
            //buffer的一半， 所以需要对最后两个分块数据的起始地址做处理
            if ((progress == (this->tileNum * BUFFER_NUM - 2)) ||
                (progress == (this->tileNum * BUFFER_NUM - 1))) {
                //分块大小变为tileLength的一半
                //倒数第2个分块数据的起始地址向前移动（tileLength-lasttileLength)，最后一个分块的起始地址以此为基础进行移动
                DataCopy(
                    inLocal[0],
                    xGm[(progress - 2) * (this->tileLength) + this->lasttileLength],
                    (this->tileLength));
                DataCopy(
                    inLocal[this->tileLength],
                    labelGm[(progress - 2) * (this->tileLength) + this->lasttileLength],
                    (this->tileLength));                
            }
            else {
                DataCopy(inLocal[0], xGm[progress * (this->tileLength)],
                        (this->tileLength));
                DataCopy(inLocal[this->tileLength], labelGm[progress * this->tileLength],
                        this->tileLength);                        
            }
        }

        inQueueIN.EnQue(inLocal);
    }    

  __aicore__ inline void Compute_defaut(int32_t progress, float p) {
        LocalTensor<T> inLocal = inQueueIN.DeQue<T>();
        LocalTensor<T> xLocal = inLocal;
        LocalTensor<T> yLocal = outQueueOUT.AllocTensor<T>();

        Abs(xLocal, xLocal, this->tileLength);
        Mul(yLocal, xLocal, xLocal, this->tileLength);

        outQueueOUT.EnQue<T>(yLocal);
        inQueueIN.FreeTensor(inLocal);
  }
__aicore__ inline void Compute_inf(int32_t progress, float p) {
        LocalTensor<T> inLocal = inQueueIN.DeQue<T>();
        LocalTensor<T> xLocal = inLocal;
        LocalTensor<T> yLocal = outQueueOUT.AllocTensor<T>();

        Abs(yLocal, xLocal, this->tileLength);

        outQueueOUT.EnQue<T>(yLocal);
        inQueueIN.FreeTensor(inLocal);
  }
__aicore__ inline void Compute_0(int32_t progress, float p) {
        LocalTensor<T> inLocal = inQueueIN.DeQue<T>();
        LocalTensor<T> xLocal = inLocal;
        LocalTensor<T> yLocal = outQueueOUT.AllocTensor<T>();

        DataCopy(yLocal, xLocal, this->tileLength);

        outQueueOUT.EnQue<T>(yLocal);
        inQueueIN.FreeTensor(inLocal);
  }
// __aicore__ inline void Compute_else(int32_t progress, float p) {
//         LocalTensor<T> inLocal = inQueueIN.DeQue<T>();
//         LocalTensor<T> xLocal = inLocal;
//         LocalTensor<T> yLocal = outQueueOUT.AllocTensor<T>();

//         Abs(xLocal, xLocal, this->tileLength);
//         LocalTensor<int32_t> tmp = tmpQ.AllocTensor<int32_t>();
//         Cast(tmp, xLocal, RoundMode::CAST_NONE, this->tileLength);
//         int16_t scalar = (int16_t)p;
//         ShiftLeft(tmp, tmp, scalar, this->tileLength);
//         if(sizeof(T)==4){
//             Cast(tmp, tmp, RoundMode::CAST_NONE, this->tileLength);
//             Cast(yLocal, tmp, RoundMode::CAST_NONE, this->tileLength);
//         }else{
//             Cast(yLocal, tmp, RoundMode::CAST_NONE, this->tileLength);
//         }
//         outQueueOUT.EnQue<T>(yLocal);
//         inQueueIN.FreeTensor(inLocal);
//   }

  __aicore__ inline void CopyOut(int32_t progress) {
    LocalTensor<T> outLocal = outQueueOUT.DeQue<T>();
    if (BUFFER_NUM == 1) {
      if (progress == this->tileNum - 1) {
        if (progress == 0) {
          //如果只有一包，则搬运的起始地址为0，tileLength为实际分块的数据量
          DataCopy(yGm[0], outLocal, this->tileLength);
        } else {
          //将最后一个分块的起始地址向前移动tileLength-lasttileLength
          DataCopy(
              yGm[(progress - 1) * this->tileLength + this->lasttileLength],
              outLocal, this->tileLength);
        }
      } else {
        DataCopy(yGm[progress * this->tileLength], outLocal,
                 this->tileLength);
      }
    }
    if (BUFFER_NUM == 2) {
      //开启double
      //buffer时，由于将输入数据分成了相等的2部分，分块大小为不开启double
      //buffer的一半， 所以需要对最后两个分块数据的起始地址做处理
      if ((progress == (this->tileNum * BUFFER_NUM - 2)) ||
          (progress == (this->tileNum * BUFFER_NUM - 1))) {
        //分块大小变为tileLength的一半
        //倒数第2个分块数据的起始地址向前移动（tileLength-lasttileLength)，最后一个分块的起始地址以此为基础进行移动
        DataCopy(
            yGm[(progress - 2) * (this->tileLength) + this->lasttileLength],
            outLocal, (this->tileLength));
      }

      else {
        DataCopy(yGm[progress * (this->tileLength)], outLocal,
                 (this->tileLength));
      }
    }

    outQueueOUT.FreeTensor(outLocal);
  }

  __aicore__ inline void FinalComput(int32_t progress, float p, float value)
  {
            LocalTensor<T> outLocal = outQueueOUT.AllocTensor<T>();
            if(p == (float)2.0) 
            {
                Duplicate(outLocal, (T)value, this->tileLength);
                Sqrt(outLocal, outLocal, this->tileLength);
            }
            // else if(p == positiveInfinity)
            // {
            //     Duplicate(outLocal, (T)value, this->tileLength);
            // }
            // else if(p == negativeInfinity)
            // {
            //     Duplicate(outLocal, (T)value, this->tileLength);
            // }
            else if(p ==(float)0.0)
            {
                Duplicate(outLocal, (T)value, this->tileLength);
            }
            else if(p == (float)1.0)
            {
                Duplicate(outLocal, (T)value, this->tileLength);
            }

            /////////////
            if (BUFFER_NUM == 1) {
            if (progress == this->tileNum - 1) {
                    if (progress == 0) {
                    //如果只有一包，则搬运的起始地址为0，tileLength为实际分块的数据量
                    DataCopy(yGm[0], outLocal, this->tileLength);
                    } else {
                    //将最后一个分块的起始地址向前移动tileLength-lasttileLength
                    DataCopy(
                        yGm[(progress - 1) * this->tileLength + this->lasttileLength],
                        outLocal, this->tileLength);
                    }
                } else {
                    DataCopy(yGm[progress * this->tileLength], outLocal,
                            this->tileLength);
                }
            }
            if (BUFFER_NUM == 2) {
                //开启double
                //buffer时，由于将输入数据分成了相等的2部分，分块大小为不开启double
                //buffer的一半， 所以需要对最后两个分块数据的起始地址做处理
                if ((progress == (this->tileNum * BUFFER_NUM - 2)) ||
                    (progress == (this->tileNum * BUFFER_NUM - 1))) {
                    //分块大小变为tileLength的一半
                    //倒数第2个分块数据的起始地址向前移动（tileLength-lasttileLength)，最后一个分块的起始地址以此为基础进行移动
                    DataCopy(
                        yGm[(progress - 2) * (this->tileLength) + this->lasttileLength],
                        outLocal, (this->tileLength));
                }

                else {
                    DataCopy(yGm[progress * (this->tileLength)], outLocal,
                            (this->tileLength));
                }
            }
            outQueueOUT.FreeTensor(outLocal);
  } 

private:
    TPipe pipe;
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueIN;
    TQue<QuePosition::VECIN, 1> inQueueNow;
    TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueOUT;
    TQue<QuePosition::VECOUT, 1> workQueue;

    TBuf<QuePosition::VECCALC> inQueueC, tmpQ;
    GlobalTensor<T> xGm;
    GlobalTensor<T> labelGm;
    GlobalTensor<T> yGm;
    uint32_t blockLength;
    uint32_t tileNum;
    uint32_t tileLength;
    uint32_t lasttileLength;
    uint32_t formerNum;
    uint32_t formerLength;
    uint32_t formertileNum;
    uint32_t formertileLength;
    uint32_t formerlasttileLength;
    uint32_t tailNum;
    uint32_t tailLength;
    uint32_t tailtileNum;
    uint32_t tailtileLength;
    uint32_t taillasttileLength;
    int length;
    float p;
};


extern "C" __global__ __aicore__ void lp_norm_v2(GM_ADDR x, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling) {
    GET_TILING_DATA(tiling_data, tiling);
    // TODO: user kernel impl
    if(tiling_data.datatype == 1){ //fp16 half 
        KernelLpnormV2<half> op;
        uint32_t tilingKey = 1;
        if (TILING_KEY_IS(1)) {
            tilingKey = 1;
        } else if (TILING_KEY_IS(2)) {
            tilingKey = 2;
        } else {
            tilingKey = 1;
        }
        op.Init(x, y, tiling_data.blockLength,
                tiling_data.tileNum, tiling_data.tileLength,
                tiling_data.lasttileLength, tiling_data.formerNum,
                tiling_data.formerLength, tiling_data.formertileNum,
                tiling_data.formertileLength, tiling_data.formerlasttileLength,
                tiling_data.tailNum, tiling_data.tailLength, tiling_data.tailtileNum,
                tiling_data.tailtileLength, tiling_data.taillasttileLength,
                tilingKey, tiling_data.size, tiling_data.p);
        op.Process(tiling_data.p);
    }else if(tiling_data.datatype == 0){ //fp32
        KernelLpnormV2<float> op;
        uint32_t tilingKey = 1;
        if (TILING_KEY_IS(1)) {
            tilingKey = 1;
        } else if (TILING_KEY_IS(2)) {
            tilingKey = 2;
        } else {
            tilingKey = 1;
        }
        op.Init(x, y, tiling_data.blockLength,
                tiling_data.tileNum, tiling_data.tileLength,
                tiling_data.lasttileLength, tiling_data.formerNum,
                tiling_data.formerLength, tiling_data.formertileNum,
                tiling_data.formertileLength, tiling_data.formerlasttileLength,
                tiling_data.tailNum, tiling_data.tailLength, tiling_data.tailtileNum,
                tiling_data.tailtileLength, tiling_data.taillasttileLength,
                tilingKey, tiling_data.size, tiling_data.p);
        op.Process(tiling_data. p);
    }  
}

